# 更新后的完整代码
from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain.tools import Tool
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings  # 更新后的导入
from langchain.docstore.document import Document
import os
from dotenv import load_dotenv
import re
import json
from datetime import datetime
import sentence_transformers  # 添加导入以管理模型缓存

# 加载环境变量
load_dotenv()
DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
LOCAL_MODEL_PATH = os.getenv("LOCAL_MODEL_PATH")
# FAISS 存储目录
FAISS_STORE_DIR = "faiss_store"
os.makedirs(FAISS_STORE_DIR, exist_ok=True)

# 设置模型缓存路径（可选）
MODEL_CACHE_PATH = os.path.expanduser("~/.cache/sentence_transformers")
os.makedirs(MODEL_CACHE_PATH, exist_ok=True)

# 初始化嵌入模型
embeddings = HuggingFaceEmbeddings(
    model_name=LOCAL_MODEL_PATH,
    cache_folder=MODEL_CACHE_PATH,  # 指定缓存路径
    model_kwargs={'device': 'cpu'},
    encode_kwargs={'normalize_embeddings': True}
)


# 验证API密钥格式的函数
def validate_api_key(api_key):
    """确保API密钥格式正确"""
    if not api_key.startswith("sk-") or len(api_key) < 30:
        raise ValueError("无效的API密钥格式。请从DeepSeek平台获取正确的密钥。")
    return True


# 符合API规范的工具名称验证
def validate_tool_name(name):
    """确保工具名称符合API规范"""
    pattern = r'^[a-zA-Z0-9_-]+$'
    if not re.match(pattern, name):
        raise ValueError(f"工具名称 '{name}' 不符合规范。只允许字母、数字、下划线和短横线。")
    return name


# 1. 创建符合规范的自定义工具函数
def ai_trend_analyzer(query: str) -> str:
    """分析AI行业趋势的专业工具"""
    # 实际应用中这里可以调用API或数据库
    return f"AI行业趋势分析: 大模型和多模态技术快速发展，特别是在{query}领域。2025年全球AI市场预计增长25%。"


def tech_forecast(query: str) -> str:
    """预测技术发展路径"""
    return f"技术预测: {query}领域将在未来2-3年内实现商业化突破。"


# 2. 初始化大模型
llm = ChatOpenAI(
    model="deepseek-chat",
    api_key=DEEPSEEK_API_KEY,
    base_url="https://api.deepseek.com/v1",
    temperature=0.7,
    max_tokens=2000
)

# 3. 创建符合规范的工具集
tools = [
    Tool(
        name=validate_tool_name("ai_trend_analyzer"),
        func=ai_trend_analyzer,
        description="Analyze latest trends in AI industry. Input should be a specific domain or technology."
    ),
    Tool(
        name=validate_tool_name("tech_forecaster"),
        func=tech_forecast,
        description="Predict development timeline for specific technologies. Input: technology name."
    )
]

# 4. 定义智能体提示词
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是DeepSeek AI行业分析师助手，可以使用工具获取专业分析。"
               "当用户询问AI行业趋势时，使用'ai_trend_analyzer'工具。"
               "当用户询问技术预测时，使用'tech_forecaster'工具。"),
    ("user", "{input}"),
    ("placeholder", "{agent_scratchpad}")
])

# 5. 创建智能体
agent = create_openai_tools_agent(
    llm=llm,
    tools=tools,
    prompt=prompt
)

# 6. 创建执行器
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    handle_parsing_errors=True,
    max_iterations=3
)


# 7. FAISS 操作函数
def init_faiss_store():
    """初始化或加载FAISS向量数据库"""
    index_path = os.path.join(FAISS_STORE_DIR, "index.faiss")
    meta_path = os.path.join(FAISS_STORE_DIR, "metadata.json")

    if os.path.exists(index_path) and os.path.exists(meta_path):
        print("加载现有的FAISS向量数据库...")
        vector_store = FAISS.load_local(
            folder_path=FAISS_STORE_DIR,
            embeddings=embeddings,
            allow_dangerous_deserialization=True
        )

        # 打印所有历史对话
        print("\n历史对话记录:")
        all_docs = vector_store.docstore._dict.values()
        for i, doc in enumerate(all_docs):
            # 跳过初始文档
            if doc.metadata.get('source') != 'init':
                print(f"对话记录 {i + 1}:")
                print(doc.page_content)
                print(f"时间: {doc.metadata.get('timestamp')}")
                print("-" * 50)
        return vector_store

    print("创建新的FAISS向量数据库...")
    return FAISS.from_texts(
        texts=["初始文档"],
        embedding=embeddings,
        metadatas=[{"source": "init", "timestamp": datetime.now().isoformat()}]
    )


def save_to_faiss(vector_store, user_input, agent_output):
    """将对话保存到FAISS向量数据库"""
    doc = Document(
        page_content=f"用户输入: {user_input}\n智能体输出: {agent_output}",
        metadata={
            "source": "agent_conversation",
            "timestamp": datetime.now().isoformat(),
            "user_input": user_input,
            "agent_output": agent_output
        }
    )

    vector_store.add_documents([doc])
    vector_store.save_local(FAISS_STORE_DIR)
    print(f"对话已保存到FAISS存储: {FAISS_STORE_DIR}")


def search_in_faiss(vector_store, query, k=3):
    """在FAISS中搜索相关对话"""
    results = vector_store.similarity_search(query, k=k)
    print(f"\n在FAISS中找到 {len(results)} 条相关对话:")
    for i, doc in enumerate(results):
        print(f"结果 {i + 1}:")
        print(f"- 内容: {doc.page_content[:150]}...")
        print(f"- 时间: {doc.metadata.get('timestamp', '未知')}")
        print("-" * 50)
    return results


# 8. 主函数
def main():
    try:
        validate_api_key(DEEPSEEK_API_KEY)
        print("正在初始化嵌入模型...")
        print("首次运行需要下载模型，可能需要几分钟，请耐心等待...")

        vector_store = init_faiss_store()
        print("嵌入模型初始化完成!")

        while True:
            # 从控制台获取用户输入
            user_input = input("\n请输入您的问题（输入'exit'退出）: ").strip()

            if user_input.lower() == 'exit':
                print("程序已退出")
                break

            # 询问是否调用DeepSeek
            use_deepseek = input("是否调用DeepSeek API进行分析? (y/n): ").strip().lower()

            if use_deepseek == 'y':
                print("\n运行智能体分析...")
                response = agent_executor.invoke({"input": user_input})
                agent_output = response["output"]

                print("\n智能体分析结果:")
                print(agent_output)

                # 保存到FAISS
                save_to_faiss(vector_store, user_input, agent_output)
            else:
                # 不调用API，直接使用本地工具
                print("\n使用本地工具进行分析...")
                # 根据问题类型选择合适的工具
                if "趋势" in user_input or "分析" in user_input:
                    agent_output = ai_trend_analyzer(user_input)
                elif "预测" in user_input or "时间表" in user_input:
                    agent_output = tech_forecast(user_input)
                else:
                    agent_output = "无法确定问题类型，请尝试包含'趋势'或'预测'等关键词"

                print("\n本地分析结果:")
                print(agent_output)

                # 保存到FAISS
                save_to_faiss(vector_store, user_input, agent_output)

            # 在FAISS中搜索相关对话
            search_in_faiss(vector_store, user_input, k=2)

    except ValueError as ve:
        print(f"验证错误: {str(ve)}")
    except Exception as e:
        print(f"执行错误: {str(e)}")


if __name__ == "__main__":
    main()